//
// Created by corey on 2021/4/13.
//

#ifndef CORNETNET_LITE_TRT_CORNER_POOL_PLUGIN_H
#define CORNETNET_LITE_TRT_CORNER_POOL_PLUGIN_H

#include "NvInfer.h"
#include "cuda_runtime.h"
#include <iostream>
#include <cstring>
#include <assert.h>

class CornerPoolPlugin : public nvinfer1::IPluginV2IOExt
{
public:
    CornerPoolPlugin()
    {

    }

    CornerPoolPlugin(const void *buffer, size_t length)
    {
        std::memcpy(&m, buffer, sizeof(m));
    }

    size_t getSerializationSize() const TRTNOEXCEPT override
    {
        return sizeof(m);
    }

    void serialize(void *buffer) const TRTNOEXCEPT override
    {
        std::memcpy(buffer, &m, sizeof(m));
    }

    IPluginV2Ext *clone() const TRTNOEXCEPT override
    {
        return new CornerPoolPlugin(&m, sizeof(m));
    }

    bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc *inOut, int nbInputs,
                                   int nbOutputs) const TRTNOEXCEPT override
    {
        if (pos < 0 || pos >= 2)
        {
            std::cerr << "only support one input and ont output\n";
            exit(0);
        }

        bool ret = false;
        switch (pos)
        {
            case 0:
                return (inOut[0].type == nvinfer1::DataType::kFLOAT || inOut[0].type == nvinfer1::DataType::kHALF || inOut[0].type == nvinfer1::DataType::kINT8) &&
                       inOut[0].format == nvinfer1::TensorFormat::kLINEAR;
            case 1:
                return inOut[pos].type == inOut[0].type && inOut[pos].format == inOut[0].format;
        }

        return ret;
    }

    int getNbOutputs() const TRTNOEXCEPT override
    { return 1; }

    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims *inputs, int nbInputDims) TRTNOEXCEPT override
    { return inputs[0]; }

    nvinfer1::DataType
    getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const TRTNOEXCEPT override
    {
        return inputTypes[0];
    }

    void configurePlugin(const nvinfer1::PluginTensorDesc *in, int nbInput, const nvinfer1::PluginTensorDesc *out,
                         int nbOutput) TRTNOEXCEPT override
    {
        m.data_type = in[0].type;
        m.input_dim = in[0].dims;
        m.format = in[0].format;
        m.scale = in[0].scale;
    }

    size_t getWorkspaceSize(int maxBatchSize) const TRTNOEXCEPT override
    { return 0; }

    int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace,
                cudaStream_t stream) TRTNOEXCEPT override
    {
        printf("CornerPool Not Implement!\n");
        return 0;
    }

    int initialize() override
    { return 0; }

    void terminate() override
    {}

    void destroy() override
    { delete this; }

    void setPluginNamespace(const char *szNamespace) override
    {}

    const char *getPluginNamespace() const override
    { return ""; }

    const char *getPluginType() const override
    { return "CornerPool"; }

    const char *getPluginVersion() const override
    { return "1"; }

    bool canBroadcastInputAcrossBatch(int inputIndex) const override
    { return false; }

    bool isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const override
    { return false; }

    void attachToContext(cudnnContext * /*cudnn*/, cublasContext * /*cublas*/,
                         nvinfer1::IGpuAllocator * /*allocator*/) override
    {}

    void detachFromContext() override
    {}

protected:
    struct
    {
        nvinfer1::DataType data_type;
        nvinfer1::TensorFormat format;
        nvinfer1::Dims input_dim;
        float scale;
    } m;
};

class CornerPoolPluginCreator : public nvinfer1::IPluginCreator
{
public:
    static nvinfer1::PluginFieldCollection plugin_field_collection;

    nvinfer1::IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override
    {
        return new CornerPoolPlugin(serialData, serialLength);
    }

    const char *getPluginName() const override
    { return "CornerPool"; }

    const char *getPluginVersion() const override
    { return "1"; }

    void setPluginNamespace(const char *szNamespace) override
    {}

    const char *getPluginNamespace() const override
    { return ""; }

    const nvinfer1::PluginFieldCollection *getFieldNames() override
    {
        return &plugin_field_collection;
    }

    nvinfer1::IPluginV2 *createPlugin(const char *name, const nvinfer1::PluginFieldCollection *fc) override
    {
        return new CornerPoolPlugin();
    }
};

void LeftPooling(const void *const pInput,
                 void *pOut,
                 const int64_t batch,
                 const int64_t channel,
                 const int64_t height,
                 const int64_t width,
                 const nvinfer1::DataType type,
                 cudaStream_t stream);

void RightPooling(const void *const pInput,
                  void *pOut,
                  const int64_t batch,
                  const int64_t channel,
                  const int64_t height,
                  const int64_t width,
                  const nvinfer1::DataType type,
                  cudaStream_t stream);

void TopPooling(const void *const pInput,
                void *pOut,
                const int64_t batch,
                const int64_t channel,
                const int64_t height,
                const int64_t width,
                const nvinfer1::DataType type,
                cudaStream_t stream);

void BottomPooling(const void *const pInput,
                   void *pOut,
                   const int64_t batch,
                   const int64_t channel,
                   const int64_t height,
                   const int64_t width,
                   const nvinfer1::DataType type,
                   cudaStream_t stream);

#endif //CORNETNET_LITE_TRT_CORNER_POOL_PLUGIN_H
